import json
import re

from llm import Embedding, enc
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


class SampleGenerator:
    def __init__(self, samples, func):
        self.samples = samples
        self.func = func
        self.embedding_matrix = []
        self.transformer_data()
        self.embedding_sample()

    def transformer_data(self):
        # 处理后应该有 embedding_content 和 return_content
        self.samples = [self.func(i) for i in self.samples]

    def embedding_sample(self):
        embedding_content_list = [i['embedding_content'] for i in self.samples]

        self.embedding_matrix = []

        batch_size = 64
        for i in range(0, len(embedding_content_list), batch_size):
            batch_content_list = embedding_content_list[i:i + batch_size]
            self.embedding_matrix.extend(Embedding(batch_content_list))

    def get_sample(self, content, top_n=3):
        embedding = Embedding(content)
        embedding = np.array(embedding).reshape(1, -1)
        similarities = cosine_similarity(embedding, self.embedding_matrix)
        top_indices = similarities.argsort()[0][-top_n:][::-1]
        return [self.samples[i]['return_content'] for i in top_indices[:top_n]]

    def add_sample(self, sample):
        sample = self.func(sample)
        self.samples.append(sample)
        self.embedding_matrix.append(Embedding(sample))


def func(x):
    return x


def get_exp_sample():
    li = []
    for path in ['submit/thought_exp.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)
                # print(data['exp'])
                if 'class ' not in data['exp']:
                    if '<thought>' in data['exp']:
                        li.append({"embedding_content": data['question'],
                                   "return_content": '<thought>' + data['exp'].split('<thought>', 1)[-1]})
                    else:
                        if len(enc.encode(data['exp'])) < 350:
                            li.append({"embedding_content": data['question'],
                                       "return_content": data['exp']})
    print('exp1 num:', len(li))
    return SampleGenerator(li, func)


def get_exp_sample2():
    li = []
    for path in ['submit/thought_exp_new.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)
                # print(data['exp'])
                if 'class ' not in data['exp']:
                    if '<thought>' in data['exp']:
                        li.append({"embedding_content": data['question'],
                                   "return_content": '<thought>' + data['exp'].split('<thought>', 1)[-1]})
                    else:
                        if len(enc.encode(data['exp'])) < 350:
                            li.append({"embedding_content": data['question'],
                                       "return_content": data['exp']})
    print('exp2 num:', len(li))
    return SampleGenerator(li, func)

from SAMPLE import SAMPLE


def get_template_sample():
    li = []
    for k, v in SAMPLE.items():
        li.append({"embedding_content": k,
                   "return_content": v})
    print('template num:', len(li))
    return SampleGenerator(li, func)


from utils.utils import grouped_dict


def get_pos_doc(text):
    tmp = []
    li = []
    for k, v in grouped_dict.items():
        for p in v:
            if k in text and p['column_name'] in text:
                if p['column_description'] not in tmp:
                    li.append({'column_name_en': p['column_name'], "column_name_zh": p['column_description']})
                    tmp.append(p['column_description'])
    return li


def get_column_sample():
    li = []
    for path in ['submit/thought_exp.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)

                if 'class ' not in data['exp']:
                    if '```md' in data['exp']:
                        li.append({"embedding_content": data['question'],
                                   "return_content": {
                                       "question": data['question'],
                                       "columns": get_pos_doc(data['exp'].split('```md', 1)[-1])}
                                   })
    print('column sample num:', len(li))
    return SampleGenerator(li, func)


from collections import defaultdict

def get_func_sample():
    li = []
    func_name_lst = []
    func_count_dict = defaultdict(int)
    for path in ['submit/thought_exp_func.jsonl']:
        with open(path, encoding='utf8') as f:
            for i in f:
                data = json.loads(i)
                func_str = re.search('```python(.*?)```', data['func'], re.DOTALL).group(1)
                func_str_list = func_str.split('if __name__ == "__main__":')
                func_str = func_str_list[0] + '\n'.join(['# ' + i[4:] for i in func_str_list[1].split('\n')])

                func_name = re.search('def (.*?)\(', func_str).group(1)
                # if func_name+"()" not in func_str:
                func_count_dict[func_name] += 1
                if func_count_dict[func_name]>1:
                    func_str = func_str.replace(f"def {func_name}", f"def {func_name}_{func_count_dict[func_name]}")
                    func_name = f"def {func_name}_{func_count_dict[func_name]}"
                li.append(
                    {"embedding_content": data['question'], 'return_content': func_str, 'func_name': func_name})

                func_name_lst.append(func_name)

    print('func sample num:', len(li))
    print('func name lst:', len(func_name_lst), len(set(func_name_lst)))
    return SampleGenerator(li, func)


if __name__ == '__main__':
    get_func_sample()
